{
 "cells": [
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-13T10:42:35.789010Z",
     "start_time": "2025-05-13T10:42:33.374610Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from d2l import torch as d2l"
   ],
   "id": "170abda05115ab86",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-13T10:42:35.819535Z",
     "start_time": "2025-05-13T10:42:35.804318Z"
    }
   },
   "cell_type": "code",
   "source": [
    "#@save\n",
    "class AttentionDecoder(d2l.Decoder):\n",
    "    \"\"\"带有注意力机制解码器的基本接口\"\"\"\n",
    "    def __init__(self, **kwargs):\n",
    "        super(AttentionDecoder, self).__init__(**kwargs)\n",
    "\n",
    "    @property\n",
    "    def attention_weights(self):\n",
    "        raise NotImplementedError"
   ],
   "id": "f05abc86b0de707d",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-13T10:42:35.866401Z",
     "start_time": "2025-05-13T10:42:35.836641Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Seq2SeqAttentionDecoder(AttentionDecoder):\n",
    "    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,\n",
    "                 dropout=0, **kwargs):\n",
    "        super(Seq2SeqAttentionDecoder, self).__init__(**kwargs)\n",
    "        self.attention = d2l.AdditiveAttention(\n",
    "            num_hiddens, num_hiddens, num_hiddens, dropout)\n",
    "        self.embedding = nn.Embedding(vocab_size, embed_size)\n",
    "        self.rnn = nn.GRU(\n",
    "            embed_size + num_hiddens, num_hiddens, num_layers,\n",
    "            dropout=dropout)\n",
    "        self.dense = nn.Linear(num_hiddens, vocab_size)\n",
    "\n",
    "    def init_state(self, enc_outputs, enc_valid_lens, *args):\n",
    "        # outputs的形状为(batch_size，num_steps，num_hiddens).\n",
    "        # hidden_state的形状为(num_layers，batch_size，num_hiddens)\n",
    "        outputs, hidden_state = enc_outputs\n",
    "        return (outputs.permute(1, 0, 2), hidden_state, enc_valid_lens)\n",
    "\n",
    "    def forward(self, X, state):\n",
    "        # enc_outputs的形状为(batch_size,num_steps,num_hiddens).\n",
    "        # hidden_state的形状为(num_layers,batch_size,\n",
    "        # num_hiddens)\n",
    "        enc_outputs, hidden_state, enc_valid_lens = state\n",
    "        # 输出X的形状为(num_steps,batch_size,embed_size)\n",
    "        X = self.embedding(X).permute(1, 0, 2)\n",
    "        outputs, self._attention_weights = [], []\n",
    "        for x in X:\n",
    "            # query的形状为(batch_size,1,num_hiddens)\n",
    "            query = torch.unsqueeze(hidden_state[-1], dim=1)\n",
    "            # context的形状为(batch_size,1,num_hiddens)\n",
    "            context = self.attention(\n",
    "                query, enc_outputs, enc_outputs, enc_valid_lens)\n",
    "            # 在特征维度上连结\n",
    "            x = torch.cat((context, torch.unsqueeze(x, dim=1)), dim=-1)\n",
    "            # 将x变形为(1,batch_size,embed_size+num_hiddens)\n",
    "            out, hidden_state = self.rnn(x.permute(1, 0, 2), hidden_state)\n",
    "            outputs.append(out)\n",
    "            self._attention_weights.append(self.attention.attention_weights)\n",
    "        # 全连接层变换后，outputs的形状为\n",
    "        # (num_steps,batch_size,vocab_size)\n",
    "        outputs = self.dense(torch.cat(outputs, dim=0))\n",
    "        return outputs.permute(1, 0, 2), [enc_outputs, hidden_state,\n",
    "                                          enc_valid_lens]\n",
    "\n",
    "    @property\n",
    "    def attention_weights(self):\n",
    "        return self._attention_weights"
   ],
   "id": "802413402a3809ca",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-13T10:42:35.959009Z",
     "start_time": "2025-05-13T10:42:35.883537Z"
    }
   },
   "cell_type": "code",
   "source": [
    "encoder = d2l.Seq2SeqEncoder(vocab_size=10, embed_size=8, num_hiddens=16,\n",
    "                             num_layers=2)\n",
    "encoder.eval()\n",
    "decoder = Seq2SeqAttentionDecoder(vocab_size=10, embed_size=8, num_hiddens=16,\n",
    "                                  num_layers=2)\n",
    "decoder.eval()\n",
    "X = torch.zeros((4, 7), dtype=torch.long)  # (batch_size,num_steps)\n",
    "state = decoder.init_state(encoder(X), None)\n",
    "output, state = decoder(X, state)\n",
    "output.shape, len(state), state[0].shape, len(state[1]), state[1][0].shape"
   ],
   "id": "28ebdbeaaacca4b7",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([4, 7, 10]), 3, torch.Size([4, 7, 16]), 2, torch.Size([4, 16]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-13T10:42:36.377029Z",
     "start_time": "2025-05-13T10:42:35.991824Z"
    }
   },
   "cell_type": "code",
   "source": [
    "embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.1\n",
    "batch_size, num_steps = 64, 10\n",
    "lr, num_epochs, device = 0.005, 250, d2l.try_gpu()\n",
    "\n",
    "train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)\n",
    "encoder = d2l.Seq2SeqEncoder(\n",
    "    len(src_vocab), embed_size, num_hiddens, num_layers, dropout)\n",
    "decoder = Seq2SeqAttentionDecoder(\n",
    "    len(tgt_vocab), embed_size, num_hiddens, num_layers, dropout)\n",
    "net = d2l.EncoderDecoder(encoder, decoder)\n",
    "d2l.train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)"
   ],
   "id": "99ed601c255790fe",
   "outputs": [
    {
     "ename": "UnicodeDecodeError",
     "evalue": "'gbk' codec can't decode byte 0xaf in position 33: illegal multibyte sequence",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mUnicodeDecodeError\u001B[0m                        Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[6], line 5\u001B[0m\n\u001B[0;32m      2\u001B[0m batch_size, num_steps \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m64\u001B[39m, \u001B[38;5;241m10\u001B[39m\n\u001B[0;32m      3\u001B[0m lr, num_epochs, device \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m0.005\u001B[39m, \u001B[38;5;241m250\u001B[39m, d2l\u001B[38;5;241m.\u001B[39mtry_gpu()\n\u001B[1;32m----> 5\u001B[0m train_iter, src_vocab, tgt_vocab \u001B[38;5;241m=\u001B[39m \u001B[43md2l\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mload_data_nmt\u001B[49m\u001B[43m(\u001B[49m\u001B[43mbatch_size\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mnum_steps\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m      6\u001B[0m encoder \u001B[38;5;241m=\u001B[39m d2l\u001B[38;5;241m.\u001B[39mSeq2SeqEncoder(\n\u001B[0;32m      7\u001B[0m     \u001B[38;5;28mlen\u001B[39m(src_vocab), embed_size, num_hiddens, num_layers, dropout)\n\u001B[0;32m      8\u001B[0m decoder \u001B[38;5;241m=\u001B[39m Seq2SeqAttentionDecoder(\n\u001B[0;32m      9\u001B[0m     \u001B[38;5;28mlen\u001B[39m(tgt_vocab), embed_size, num_hiddens, num_layers, dropout)\n",
      "File \u001B[1;32m~\\anaconda3\\envs\\d2l\\lib\\site-packages\\d2l\\torch.py:931\u001B[0m, in \u001B[0;36mload_data_nmt\u001B[1;34m(batch_size, num_steps, num_examples)\u001B[0m\n\u001B[0;32m    929\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[38;5;21mload_data_nmt\u001B[39m(batch_size, num_steps, num_examples\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m600\u001B[39m):\n\u001B[0;32m    930\u001B[0m \u001B[38;5;250m    \u001B[39m\u001B[38;5;124;03m\"\"\"Return the iterator and the vocabularies of the translation dataset.\"\"\"\u001B[39;00m\n\u001B[1;32m--> 931\u001B[0m     text \u001B[38;5;241m=\u001B[39m preprocess_nmt(\u001B[43mread_data_nmt\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m)\n\u001B[0;32m    932\u001B[0m     source, target \u001B[38;5;241m=\u001B[39m tokenize_nmt(text, num_examples)\n\u001B[0;32m    933\u001B[0m     src_vocab \u001B[38;5;241m=\u001B[39m d2l\u001B[38;5;241m.\u001B[39mVocab(source, min_freq\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m2\u001B[39m,\n\u001B[0;32m    934\u001B[0m                           reserved_tokens\u001B[38;5;241m=\u001B[39m[\u001B[38;5;124m'\u001B[39m\u001B[38;5;124m<pad>\u001B[39m\u001B[38;5;124m'\u001B[39m, \u001B[38;5;124m'\u001B[39m\u001B[38;5;124m<bos>\u001B[39m\u001B[38;5;124m'\u001B[39m, \u001B[38;5;124m'\u001B[39m\u001B[38;5;124m<eos>\u001B[39m\u001B[38;5;124m'\u001B[39m])\n",
      "File \u001B[1;32m~\\anaconda3\\envs\\d2l\\lib\\site-packages\\d2l\\torch.py:862\u001B[0m, in \u001B[0;36mread_data_nmt\u001B[1;34m()\u001B[0m\n\u001B[0;32m    860\u001B[0m data_dir \u001B[38;5;241m=\u001B[39m d2l\u001B[38;5;241m.\u001B[39mdownload_extract(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mfra-eng\u001B[39m\u001B[38;5;124m'\u001B[39m)\n\u001B[0;32m    861\u001B[0m \u001B[38;5;28;01mwith\u001B[39;00m \u001B[38;5;28mopen\u001B[39m(os\u001B[38;5;241m.\u001B[39mpath\u001B[38;5;241m.\u001B[39mjoin(data_dir, \u001B[38;5;124m'\u001B[39m\u001B[38;5;124mfra.txt\u001B[39m\u001B[38;5;124m'\u001B[39m), \u001B[38;5;124m'\u001B[39m\u001B[38;5;124mr\u001B[39m\u001B[38;5;124m'\u001B[39m) \u001B[38;5;28;01mas\u001B[39;00m f:\n\u001B[1;32m--> 862\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mf\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mread\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n",
      "\u001B[1;31mUnicodeDecodeError\u001B[0m: 'gbk' codec can't decode byte 0xaf in position 33: illegal multibyte sequence"
     ]
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "engs = ['go .', \"i lost .\", 'he\\'s calm .', 'i\\'m home .']\n",
    "fras = ['va !', 'j\\'ai perdu .', 'il est calme .', 'je suis chez moi .']\n",
    "for eng, fra in zip(engs, fras):\n",
    "    translation, dec_attention_weight_seq = d2l.predict_seq2seq(\n",
    "        net, eng, src_vocab, tgt_vocab, num_steps, device, True)\n",
    "    print(f'{eng} => {translation}, ',\n",
    "          f'bleu {d2l.bleu(translation, fra, k=2):.3f}')"
   ],
   "id": "1c416138ac4a86b8"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "attention_weights = torch.cat([step[0][0][0] for step in dec_attention_weight_seq], 0).reshape((\n",
    "    1, 1, -1, num_steps))"
   ],
   "id": "e08743db5566bdff"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-13T10:42:36.380677900Z",
     "start_time": "2025-05-13T10:42:18.349651Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 加上一个包含序列结束词元\n",
    "d2l.show_heatmaps(\n",
    "    attention_weights[:, :, :, :len(engs[-1].split()) + 1].cpu(),\n",
    "    xlabel='Key positions', ylabel='Query positions')"
   ],
   "id": "8db89aca6c42147",
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'd2l' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[1], line 2\u001B[0m\n\u001B[0;32m      1\u001B[0m \u001B[38;5;66;03m# 加上一个包含序列结束词元\u001B[39;00m\n\u001B[1;32m----> 2\u001B[0m \u001B[43md2l\u001B[49m\u001B[38;5;241m.\u001B[39mshow_heatmaps(\n\u001B[0;32m      3\u001B[0m     attention_weights[:, :, :, :\u001B[38;5;28mlen\u001B[39m(engs[\u001B[38;5;241m-\u001B[39m\u001B[38;5;241m1\u001B[39m]\u001B[38;5;241m.\u001B[39msplit()) \u001B[38;5;241m+\u001B[39m \u001B[38;5;241m1\u001B[39m]\u001B[38;5;241m.\u001B[39mcpu(),\n\u001B[0;32m      4\u001B[0m     xlabel\u001B[38;5;241m=\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mKey positions\u001B[39m\u001B[38;5;124m'\u001B[39m, ylabel\u001B[38;5;241m=\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mQuery positions\u001B[39m\u001B[38;5;124m'\u001B[39m)\n",
      "\u001B[1;31mNameError\u001B[0m: name 'd2l' is not defined"
     ]
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "c9b9cb96ddf46e3"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.22"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
